import pygame
import binascii
import os
import time
import socket
import sys
import threading
import struct, pickle
import asyncio
import cv2

class ServoInfo:
	def __init__(self, servo_id, servo_angle):
		self.id = servo_id
		self.angle = servo_angle
		self.init_angle = servo_angle
	def turn(self, delta_angle):
		self.angle += delta_angle
		if self.angle > 180:
			self.angle = 180
		if self.angle < 0:
			self.angle = 0
		return self.id, self.angle
	def setAngle(self, angle):
		self.angle = angle
		if self.angle > 180:
			self.angle = 180
		if self.angle < 0:
			self.angle = 0
		return self.id, self.angle
	def reSetAngle(self):
		self.angle = self.init_angle
		return self.id, self.angle

class CarController:
	def __init__(self, IP = None, rate = 1.0, speed_rate = 1.0, port = 2333, image_on = False):
		self.ip = IP
		self.port = port # 若不开启图像，则port, port+1占用，若开启则port+2也占用
		self.image_on = image_on
		self.init()
	def __del__(self):
		self.quit()
	
	def init(self):
		if self.ip is None: return None
		
		print('here at init()')
		
		try:
			self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
			self.sock.connect((self.ip, self.port))
			self.sock.send('COMMAND'.encode())
			time.sleep(.1)
		except socket.error as msg:
			print(msg)
			sys.exit(1)
		self.sock2 = self.sock
		self.distance = None
		self.angle = None
		self.horizontal_servo = None
		# try:
		# 	self.sock2 = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
		# 	self.sock2.connect((self.ip, self.port + 1))
		# 	self.sock2.send('GET_DISTANCE'.encode())
		# 	self.distance = None
		# 	self.angle = None
		# 	self.horizontal_servo = None
		# 	time.sleep(.1)
		# except socket.error as msg:
		# 	print(msg)
		# 	sys.exit(1)
		
		if self.image_on:
			try:
				self.sock_image = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
				self.sock_image.connect((self.ip, self.port + 2))
				self.sock_image.send('IMAGE'.encode())
				time.sleep(.1)
			except socket.error as msg:
				print(msg)
				sys.exit(1)
		else: self.sock_image = None
		
		return None
	def quit(self):
		if self.ip is not None:
			self.sendMsgCharByChar('\xff\xef\xef\xee\xff')
			self.sock.close()
			self.sock2.close()
			if self.sock_image is not None:
				self.sock_image.send('exit'.encode())
				self.sock_image.close()
		return None
	
	def sendMsgCharByChar(self, msg):
		msg = '\xFF' + msg + '\xFF'
		for c in msg:
			self.sock.send(struct.pack('B', ord(c)))
			time.sleep(.005)
		return None

	def motorStop(self):
		self.sendMsgCharByChar('\x00\x00\x00')
		return None
	def motorForward(self):
		self.sendMsgCharByChar('\x00\x01\x00')
		return None
	def motorBackward(self):
		self.sendMsgCharByChar('\x00\x02\x00')
		return None
	def motorLeft(self):
		self.sendMsgCharByChar('\x00\x03\x00')
		return None
	def motorRight(self):
		self.sendMsgCharByChar('\x00\x04\x00')
		return None
	def motorLeftForward(self):
		self.sendMsgCharByChar('\x00\x05\x00')
		return None
	def motorLeftBackward(self):
		self.sendMsgCharByChar('\x00\x06\x00')
		return None
	def motorRightForward(self):
		self.sendMsgCharByChar('\x00\x07\x00')
		return None
	def motorRightBackward(self):
		self.sendMsgCharByChar('\x00\x08\x00')
		return None
	def servoTurnRight(self, angle = 10):
		if self.horizontal_servo == None:
			return None
		self.horizontal_servo.turn(-angle)
		self.sendMsgCharByChar('\x01' + chr(self.horizontal_servo.id) + chr(self.horizontal_servo.angle))
		return None
	def servoTurnLeft(self, angle = 10):
		if self.horizontal_servo == None:
			return None
		self.horizontal_servo.turn(angle)
		self.sendMsgCharByChar('\x01' + chr(self.horizontal_servo.id) + chr(self.horizontal_servo.angle))
		return None
	def servoReSet(self):
		if self.horizontal_servo == None:
			return None
		self.horizontal_servo.reSetAngle()
		self.sendMsgCharByChar('\x01' + chr(self.horizontal_servo.id) + chr(self.horizontal_servo.angle))
		return None
	
	def receiveMessage(self):
		now_data, receiving = [], False
		while True:
			try:
				data = binascii.b2a_hex(self.sock2.recv(1)).decode()
			except socket.error as msg:
				print('receiveCommand():', msg)
				break
			if data == 'ff':
				if receiving and len(now_data) == 3:
					receiving = False
					self.distance = int(now_data[1], 16)
					self.angle = int(now_data[2], 16)
					if self.horizontal_servo is None:
						self.horizontal_servo = ServoInfo(7, self.angle)
				else:
					receiving = True
					now_data = []
			else:
				if receiving: now_data.append(data)
		return None
	def getMessageOnce(self): # 获取一次测距数据，返回元组：(距离值, 角度值)
		if self.horizontal_servo is None:
			return None
		return self.distance, self.angle
	
	def getImage(self):
		if self.sock_image is None: return None
		recv_data = self.sock_image.recv(8)
		print('CarController::recv:', recv_data)
		s_type, size = struct.unpack('4si', recv_data)
		print('CarController::recv unpack', s_type, size)
		
		res = b''
		if size <= 0 or s_type != b'SIZE':
			print('Image fail.')
			self.sock_image.send('ERR1'.encode())
			return None
		self.sock_image.send(s_type)
		while True:
			if size - len(res) <= 146000:
				data = self.sock_image.recv(size - len(res))
				# print('CarController::', len(data))
				res += data
				break
			data = self.sock_image.recv(146000)
			# print('CarController::', len(data))
			res += data
			pass
		if len(res) != size:
			print('error')
			self.sock_image.send('ERR2'.encode())
			return None
		self.sock_image.send('PIC'.encode())
		
		try:
			res = pickle.loads(res)
		except:
			return None
		
		res = res.transpose(1, 0, 2)[:,:,::-1]
		# print('CarController::img_success!')
		return res


if __name__ == '__main__': # test
	car_controller = CarController(IP=input("ip?"), image_on=False)
	car_controller.servoTurnRight(20)
	car_controller.servoTurnLeft()
	car_controller.servoReSet()